import json
import asyncio
import time
from redis.asyncio import Redis
from config.env import RedisConfig, AppConfig
from ocr_core.schema import OcrTask, OcrTaskResult, TaskStatus, LlmTask
from ocr_core.llm_producer import LlmProducer
from paddleocr import PaddleOCR
from utils.ocr_llm_util import ImageOptimizer
from utils.log_util import logger
from utils.upload_util import UploadUtil
from config.mongodb import MongoDB, OcrResultMDB, OcrTaskMDB, TaskMDB


class OcrConsumer:
    def __init__(
        self,
        redis_client: Redis,
        ocr_config: dict = None,
        concurrency: int = 8,
    ):
        self.redis = redis_client
        self.concurrency = concurrency
        self.running = True
        self.ocr_config = ocr_config or {}
        self.ocr = None
        MongoDB.init()

    async def start(self):
        """启动 N 个并发消费者"""
        await self._init_ocr()
        logger.debug(
            f"OCR 任务消费者启动，队列={RedisConfig.redis_ocr_queue}, 并发数={self.concurrency}"
        )
        workers = [
            asyncio.create_task(self._consume_loop(i)) for i in range(self.concurrency)
        ]
        await asyncio.gather(*workers)

    async def stop(self):
        """停止消费者"""
        self.running = False
        logger.debug("任务消费者停止中...")

    async def _init_ocr(self):
        """初始化 OCR 模型"""
        self.ocr = PaddleOCR(**self.ocr_config)
        logger.debug("OCR 模型初始化完成")

    async def _consume_loop(self, worker_id: int):
        """单个 worker 的消费循环"""
        logger.debug(f"OCR Worker-{worker_id} 已启动")
        while self.running:
            try:
                task_msg = await self.redis.blpop(
                    RedisConfig.redis_ocr_queue,
                    timeout=RedisConfig.consumer_timeout_seconds,
                )
                if not task_msg:
                    continue

                _, serialized_task = task_msg
                ocr_task = OcrTask(**json.loads(serialized_task))
                await self._handle_task(ocr_task, worker_id)
                # OSS 上传文件
                find_one = OcrTaskMDB.find_one({"task_id": ocr_task.task_id})
                oss_path = UploadUtil.upload_to_oss(
                    find_one["optimized_path"],
                    find_one["optimized_path"].replace("vf_admin/download_path/", ""),
                )
                OcrTaskMDB.update_one(
                    {"task_id": ocr_task.task_id, "optimized_path": oss_path}
                )

            except (ConnectionError, TimeoutError) as e:
                logger.error(f"OCR Worker-{worker_id} Redis 连接错误: {e}，等待重试...")
                await asyncio.sleep(5)
            except Exception as e:
                logger.error(f"OCR Worker-{worker_id} 消费循环异常: {e}")
                await asyncio.sleep(1)

    async def _handle_task(self, task: OcrTask, worker_id: int):
        """处理单个任务"""
        task_key = f"{RedisConfig.redis_ocr_result}{task.task_id}"
        task_result_raw = await self.redis.get(task_key)

        if not task_result_raw:
            logger.warning(
                f"OCR Worker-{worker_id} 未找到任务 {task.task_id} 的初始状态"
            )
            return

        task_result = OcrTaskResult(**json.loads(task_result_raw))
        await self._update_task_result(
            task_result.replace(status=TaskStatus.PROCESSING)
        )

        try:
            result = await self._process_task(task)
            await self._update_task_result(result)
            logger.debug(f"OCR Worker-{worker_id} 成功完成任务 {task.task_id}")
        except Exception as e:
            logger.error(f"OCR Worker-{worker_id} 任务 {task.task_id} 失败: {e}")
            await self._update_task_result(
                task_result.replace(status=TaskStatus.FAILED, error=str(e))
            )

    async def _update_task_result(self, task_result: OcrTaskResult):
        """更新任务结果"""
        await self.redis.set(
            f"{RedisConfig.redis_ocr_result}{task_result.task_id}",
            json.dumps(task_result.model_dump(), ensure_ascii=False),
            ex=3600,
        )
        OcrResultMDB.update_one(task_result.model_dump())

    async def _process_task(self, task: OcrTask) -> OcrTaskResult:
        """OCR 任务处理逻辑"""
        time_start = time.time()
        optimize_time = 0
        report = None

        # 图片优化
        optimizer = ImageOptimizer(task.image_path)
        optimizer.optimize(quality=85, max_dimension=3840)
        optimize_time = int((time.time() - time_start) * 1000)
        report = optimizer.get_optimization_report()
        optimized_path = report["output_path"] if report else task.image_path
        OcrTaskMDB.update_one(
            {"task_id": task.task_id, "optimized_path": optimized_path}
        )
        # OCR 识别
        if AppConfig.app_env == "prod":
            result = self.ocr.predict(optimized_path)
            rec_texts = result[0]["rec_texts"]
        else:
            rec_texts = [
                "封口签",
                "YH001胶囊/安慰剂",
                "方案编号：YH001AA-1101",
                "药物编号：D0987",
                "药物提供方：杭州禹泓医药科技有限公司",
            ] * 2

        ocr_time = int((time.time() - time_start) * 1000) - optimize_time

        # 投递 LLM 队列
        if task.llm:
            llm_producer = LlmProducer(self.redis)
            await llm_producer.create_llm_task(
                LlmTask(
                    task_id=task.task_id,
                    scenario=task.scenario,
                    user_prompt=task.user_prompt,
                    ocr_text=json.dumps(rec_texts, ensure_ascii=False),
                )
            )
            logger.debug(f"任务 {task.task_id} 已投递至 LLM 队列")

        return OcrTaskResult(
            task_id=task.task_id,
            scenario=task.scenario,
            texts=rec_texts,
            img_ms=optimize_time,
            ocr_ms=ocr_time,
            llm_ms=0,
            status=TaskStatus.COMPLETED,
            error="",
        )
